import matplotlib.pyplot as plt
import numpy as np
from utils.data_to_color import ImColor_tp, ImColor_tp_12mins
from utils.diamond131 import SWAN
from scipy.interpolate import griddata
import glob
from utils.config import station_grid_path, radar_path, weight_root, label_path
import os
from makeAxData import get_APCP
from datetime import datetime, timedelta
from model.SmaAt_UNet import SmaAt_UNet
import json
from utils.tools import ts_station, ts
import torch.nn.functional as F
import torch
from utils.config import *
import random
from utils.s2d import space_to_depth2
from utils.d2s import depth_to_space

events = [104.2, 112.19, 21.01, 26.6]


def get_rad(radfile):
    a2 = 118.239
    b2 = 1.5241
    min_z = -5.
    max_z = 55.

    f = SWAN(radfile)
    data = f.get_data()

    lat, lon, rad = data.variables.values()
    rad = np.minimum(max_z, rad)
    rad = np.maximum(min_z, rad)

    rad_r = np.power(10, (rad - 10 * np.log10(a2)) / (10 * b2))
    return np.nan_to_num(rad_r), lat.values, lon.values


def compare(data_line):
    apcp_file, radar_file, station_file = data_line.strip().split(' ')
    stime = os.path.dirname(apcp_file).split('GXR1KM_APCP_')[1][:12]
    ftime = os.path.basename(apcp_file)[:-4]
    bj_ftime = os.path.basename(station_file)[:-4]
    fmins = (datetime.strptime(ftime, '%Y%m%d%H%M') - datetime.strptime(stime, '%Y%m%d%H%M')).seconds // 60
    # if int(ftime[:10]) != 2021062513:
    #     return
    print(apcp_file, fmins)

    apcp_data = np.loadtxt(apcp_file).copy()[120:-200, 200:-160]
    # apcp_data = np.loadtxt(apcp_file).copy()[120:280, 200:360]
    apcp_data[apcp_data < 0] = 0

    input_data = np.expand_dims(apcp_data, axis=0)
    # input_data = np.expand_dims(cv2.resize(apcp_data, (160, 112)), axis=0)

    raddata, lat0, lon0 = get_rad(radar_file)
    raddata[raddata < 0] = 0
    raddata = raddata[140:-200, 220:-180] / 5
    # raddata = raddata[140:300, 220:380] / 5

    x_batch = torch.from_numpy(input_data.copy()[np.newaxis, ...])
    x_batch = space_to_depth2(x_batch, 2)

    with torch.no_grad():
        out = model(x_batch.to(device=device, dtype=torch.float32))
        # out = F.max_pool2d(out, 3, 1, padding=1)
        out = depth_to_space(out, 2).cpu().numpy().squeeze()

        out[out < 0] = 0

        # out11 = F.max_pool2d(out.copy(), 3, 1, padding=1).numpy().squeeze()
        # out_r = cv2.resize(out, (ydim, xdim))

    # sub_data[apcp_data > out] = (apcp_data[apcp_data > out] + out[apcp_data > out]) * 0.5
    sub_data1 = np.maximum(apcp_data, out)
    # a = sub_data1.copy()
    # a[a < 1.5] = 0
    sub_data0 = np.where(apcp_data < 0.6, out, sub_data1)
    # sub_data0 = np.where(apcp_data > 5, sub_data1, out)
    # apcp = apcp_data.copy()
    # apcp[apcp_data > 1.3] = 0
    # sub = sub_data1.copy()
    # sub[sub<5] = 0
    # out[out > 1] *= 1.2
    # sub_data0 = np.maximum(apcp, out)
    # sub_data0[sub_data1 > 4.5] = sub_data1[sub_data1 > 4.5] * 1.2
    # sub_data[apcp_data < 1] = apcp_data[apcp_data < 1]

    station_data = np.nan_to_num(np.loadtxt(station_file))
    # station_data = np.load(station_file)[140:-200, 220:-180]

    rs_data = np.where(station_data > raddata, station_data, raddata)

    # rs_data_torch = torch.from_numpy(rs_data).unsqueeze(0).unsqueeze(0)
    # rs_data = F.max_pool2d(rs_data_torch, 11, 1, padding=5).numpy().squeeze()

    # station_data_r = cv2.resize(station_data, (160, 112))
    maex = np.abs(apcp_data - station_data).mean()
    maey = np.abs(out - station_data).mean()
    Amaex.append(maex)
    Amaey.append(maey)

    maeup = (maex - maey) / maex * 100
    tsx = ts_station(apcp_data, station_data, idxs).round(2)
    # out[out > 0] += 5
    tsy = ts_station(out, station_data, idxs).round(2)
    tss0 = ts_station(sub_data0, station_data, idxs).round(2)
    tss1 = ts_station(sub_data1, station_data, idxs).round(2)

    # tsx = ts(apcp_data, station_data).round(2)
    # tsy = ts(out, station_data).round(2)
    # tss = ts(sub_data, station_data).round(2)
    Atsx.append(tsx)
    Atsy.append(tsy)
    Atss0.append(tss0)
    Atss1.append(tss1)

    tsup = np.divide((tsy - tsx) * 100, tsx, out=np.empty_like(tsx), where=tsx != 0).round(2)
    tsup1 = np.divide((tss0 - tsx) * 100, tsx, out=np.empty_like(tsx), where=tsx != 0).round(2)
    tsup2 = np.divide((tss1 - tsx) * 100, tsx, out=np.empty_like(tsx), where=tsx != 0).round(2)

    g = np.zeros_like(tsx)
    g[tsx < tss0] = 1
    g[tsx > tss0] = -1
    good.append(g)
    print(
        # f"MAEX:{maex:.2f}, MAEY:{maey:.2f}, "
        f"MAE_UP:{maeup:.2f}% UP0:{tsup.tolist()}, UP1:{tsup1.tolist()}, UP2:{tsup2.tolist()}%")

    # print(f"MAEX:{maex:.2f}, MAEY:{maey:.2f}, UP:{maeup:.2f}% ")

    plt.suptitle(f'UTC_{stime}_{ftime}  f:{fmins}mins')

    plt.subplot(131)
    plt.imshow(ImColor_tp_12mins(apcp_data.copy()), extent=events)
    plt.title(f'apcp_data {tsx}')

    # plt.subplot(232)
    # plt.imshow(ImColor_tp_12mins(out.copy()), extent=events)
    # plt.title(f'out {tsy}')

    # plt.subplot(233)
    # plt.imshow(ImColor_tp_12mins(sub_data0.copy()), extent=events)
    # plt.title(f'sub_data0 {tss0}')

    plt.subplot(132)
    plt.imshow(ImColor_tp_12mins(sub_data1.copy()), extent=events)
    plt.title(f'revised_data {tss1}')

    #
    plt.subplot(133)
    plt.imshow(ImColor_tp_12mins(station_data.copy()), extent=events)
    plt.title('station_data')
    # #
    # plt.subplot(236)
    # plt.imshow(ImColor_tp_12mins(raddata))
    # plt.title('raddata')
    # #
    # plt.subplot(236)
    # plt.imshow(ImColor_tp_12mins(rs_data))
    # plt.title('rs_data')

    # plt.subplot(234)
    # plt.imshow(ImColor_tp(nearest * 5 * 24))
    # plt.title('nearest')
    #

    # plt.subplot(236)
    # plt.imshow(ImColor_tp_12mins(rs_max_pool.numpy().squeeze()))
    # plt.title('rs_max_pool')

    plt.gcf().set_size_inches(12, 6)
    plt.tight_layout()
    # plt.show()
    result_file = f"result/UTC_{stime}_{ftime}.jpg"
    plt.savefig(result_file)
    # plt.clf()


if __name__ == '__main__':

    with open('/home/gym/projects/my_test/data/idxs.json') as f:
        idxs = tuple(np.array(json.load(f)).T)

    xdim, ydim, reso = 560, 800, 0.01

    model = SmaAt_UNet(in_channels=4, out_channels=4)
    # model = UNet()
    device = torch.device("cuda:0")
    weight_file = os.path.join(weight_root, f"20220304.pth")
    print(weight_file)
    checkpoints = torch.load(weight_file, map_location=device)
    state_dict = checkpoints['state_dict']
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()

    txtpath = f"{txt_path}/tp_test.txt"
    txtList = open(txtpath, 'r').readlines()[200:300]
    # random.seed(222)
    # random.shuffle(txtList)

    Amaex, Amaey, Atsx, Atsy, Atss0, Atss1, good = [], [], [], [], [], [], []
    for i, data_path in enumerate(txtList):
        # if i % 60 == 0:
        compare(data_path)

    print(np.mean(Amaex).round(2), np.mean(Amaey).round(2), np.sum(good, axis=0))
    print("apcp:", np.vstack(Atsx).mean(axis=0).round(4).tolist())
    print("out :", np.vstack(Atsy).mean(axis=0).round(4).tolist())
    print("sub_data0 :", np.vstack(Atss0).mean(axis=0).round(4).tolist())
    print('sub_data1:', np.vstack(Atss1).mean(axis=0).round(4).tolist())
